#include "irfunction.h"

#include <functional>
#include <iostream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "effect.h"
#include "ir.h"
#include "irblock.h"
#include "iroperand.h"
#include "utils.h"

using namespace Utils;

#define VAR BLUE
#define OP DEF
#define LABEL COLOR(5)

void IRFunction::Simplify() {
    std::unordered_set<IRBlock*> vis;
    std::function<void(IRBlock*)> dfs1 = [&](IRBlock* cur) {
        Assert(cur->GetTerminator(), "IRBlock should end with a terminator");
        if (vis.count(cur)) return;
        vis.insert(cur);
        for (auto nxt : cur->GetNext()) {
            dfs1(nxt);
        }
    };
    dfs1(entry);

    // delete redundant reverse edges
    for (auto u : vis) {
        for (auto it = u->prev.begin(); it != u->prev.end();) {
            if (vis.count(it->first))
                ++it;
            else
                it = u->prev.erase(it);
        }
    }

    // merge block chains
    std::function<void(IRBlock*, IRBlock*)> dfs2 = [&](IRBlock* pre,
                                                       IRBlock* cur) {
        if (pre && cur->prev.size() == 1 &&
            pre->GetTerminator()->As<IR::Jump>()) {
            pre->Remove();
            pre->irList.pop_back();
            cur->irList.insert(cur->irList.begin(), pre->irList.begin(),
                               pre->irList.end());
            if (pre == entry) entry = cur;
        }
        if (vis.count(cur)) return;
        vis.insert(cur);
        if (auto jmp = cur->GetTerminator()->As<IR::Jump>()) {
            dfs2(cur, jmp->nxt);
        } else if (auto br = cur->GetTerminator()->As<IR::Branch>()) {
            dfs2(cur, br->t), dfs2(cur, br->f);
        }
        // for (auto nxt : cur->GetNext()) dfs2(cur, nxt);
    };
    vis.clear();
    dfs2(nullptr, entry);
}

void IRFunction::SetParent() {
    std::unordered_set<IRBlock*> vis;
    std::function<void(IRBlock*)> dfs = [&](IRBlock* cur) {
        if (vis.count(cur)) return;
        vis.insert(cur);
        for (auto phi : cur->phis) phi->parent = cur;
        for (auto ir : cur->irList) ir->parent = cur;
        for (auto nxt : cur->GetNext()) dfs(nxt);
    };
    dfs(entry);
}

void IRFunction::Validate() {
    return;
    std::unordered_set<IRBlock*> vis;
    std::unordered_map<IRBlock*, int> indeg;
    std::function<void(IRBlock*)> dfs = [&](IRBlock* cur) {
        if (vis.count(cur)) return;
        vis.insert(cur);
        for (auto phi : cur->phis)
            Assert(phi->parent == cur, "parent not correct");
        for (auto ir : cur->irList)
            Assert(ir->parent == cur, "parent not correct");
        for (auto nxt : cur->GetNext()) {
            Assert(nxt->prev.count(cur), "reverse edge not found");
            ++indeg[nxt], dfs(nxt);
        }
    };
    dfs(entry);

    for (auto [u, deg] : indeg) {
        Assert(u->prev.size() == deg,
               "valid checker failed:\n"
               "  indeg (%d) not equal to prev.size() (%d)",
               deg, u->prev.size());
    }
}

void IRFunction::AnalEffects() {
    globalEffects.Clear();
    usedGlobalVars.clear();
    argMemEffects.clear();

    std::unordered_map<std::string, int> argId;
    for (int i = 0; i < args.size(); ++i) argId[args[i]->name] = i;

    std::unordered_set<IRBlock*> vis;
    std::function<void(IRBlock*)> dfs = [&](IRBlock* u) {
        if (vis.count(u)) return;
        vis.insert(u);
        for (auto ir : u->irList) {
            for (auto v : ir->UsedVar()) {
                if (!v->isGlobal) continue;
                usedGlobalVars[v->name] = *v;
            }
            auto effects = ir->Effects();
            for (auto [_, v] : effects.mem) {
                if (!v.isGlobal) {
                    if (!argId.count(v.name)) continue;
                    argMemEffects.push_back(argId[v.name]);
                } else {
                    globalEffects.AddMem(&v);
                }
            }
            for (auto [_, v] : effects.var) {
                if (!v.isGlobal) continue;
                globalEffects.AddVar(&v);
            }
            globalEffects.io |= effects.io;
        }
        for (auto v : u->GetNext()) dfs(v);
    };
    dfs(entry);
}
